import os
import code2word as c2w

def identify_file_type(file_path):
    _, ext = os.path.splitext(file_path)
    return ext.lower()

def read_files_and_label(parent_folder):
    train_set = []  # 存储训练集代码片段
    labels = []  # 存储标签
    # 八个子文件夹名称列表，分别对应八项分类
    subfolders = ['normal', 'backdoor', 'dos_attack', 'ftp_exfiltration', 'keylogger', 'persistence_backdoor', 'ransomware', 'reverse_shell']

    for subfolder in subfolders:
        folder_path = os.path.join(parent_folder, subfolder)
        # 遍历子文件夹中的所有文件
        for filename in os.listdir(folder_path):
            # 构建完整的文件路径
            file_path = os.path.join(folder_path, filename)
            # 检查是否为文件
            if os.path.isfile(file_path):
                file_type = identify_file_type(file_path)
                # 读取并转换代码段
                if file_type == ".py":
                    code = c2w.py_code2word(file_path)
                elif file_type == ".java":
                    code = c2w.go_Java_code2word(file_path)
                elif file_type == ".go":
                    code = c2w.go_Java_code2word(file_path)
                else:
                    code = []
                    print("Unsupported file type\n")
                train_set.append(code)
                # 将标签转换为对应的索引
                labels.append(subfolders.index(subfolder))

    return train_set, labels

# 调用函数
if __name__ == "__main__":
    parent_folder = '../data/train_data'  # 父文件夹路径
    train_set, labels = read_files_and_label(parent_folder)
    for each in train_set:
        print(each)
        print('\n')
    print(labels)
